import numpy as np
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import skimage
from PyQt5 import QtWidgets
import anndata as ad
import matplotlib.pyplot as plt
import cv2
import skimage.transform


class Our_Image():
    def __init__(self,img_preserve):
        self.img = None
        self.img_temp = None
        self.img_preserve = img_preserve
        self.height = img_preserve.shape[0]
        self.width = img_preserve.shape[1]

        self.angle = 0
        self.move_x = 0
        self.move_y = 0
        self.scale = 1.0
        self.scale_all = 1.0
        self.flip_V = False
        self.flip_H = False
        self.a =0
        self.b =0
        self.c =0
        self.init_set()

    def init_set(self):
        if self.height>= 2500 or self.width >= 2500:
            if len(self.img_preserve.shape) <= 2:
                self.img = skimage.transform.rescale(self.img_preserve, 800/self.height)
                self.img_black_temp = self.img.copy()
            else:
                self.img = cv2.resize(self.img_preserve, (0, 0), fx=800/self.height, fy=800/self.height)
            self.scale = 800/self.height
        else:
            self.img = self.img_preserve.copy()

        if len(self.img_preserve.shape) > 2:
            self.img_black = skimage.color.rgb2gray(self.img)
        else:
            self.img_black = self.img.copy()

        self.maxValue = np.max(self.img_black.reshape(-1))

    def draw_img(self,grap):
        canvas = plt_img(self.img)
        add_to_graph(canvas, grap)
        plt.close('all')

    def change_color(self, grap_img, grap_hist, low, top):
        low = low / 300 * self.maxValue
        top = top / 300 * self.maxValue

        self.img_black_temp = set_range(self.img_black, low, top)
        canvas = plt_img(self.img_black_temp)
        add_to_graph(canvas, grap_img)
        canvas = plt_hist(self.img_black_temp)
        add_to_graph(canvas, grap_hist)
        plt.close('all')

    def rescale(self, scale, grap_img):
        if len(self.img.shape) <=2:
            self.img_temp = skimage.transform.rescale(self.img_preserve, scale)
        else:
            self.img_temp = cv2.resize(self.img_preserve,(0,0),fx=scale,fy=scale)
        self.scale = scale
        canvas = plt_img(self.img_temp)
        add_to_graph(canvas, grap_img)
        plt.close('all')
        self.save()

    def flip(self,stage,grap_img):
        if stage == 0:
            self.flip_V = not self.flip_V
        if stage == 1:
            self.flip_H = not self.flip_H
        self.img_temp = cv2.flip(self.img,stage)
        canvas = plt_img(self.img_temp)
        add_to_graph(canvas, grap_img)
        plt.close('all')
        self.save()

    def save(self):
        if len(self.img.shape) <= 2:
            #self.img_preserve = self.img_black_temp.copy()
            self.img_black = self.img_black_temp.copy()
            self.img = self.img_black_temp.copy()
            self.img_black_temp = None

        if type(self.img_temp) != type(None):
            self.img = self.img_temp.copy()
            self.img_temp = None 

    def remove_images(self):
        self.img = None
        self.img_temp = None
        self.img_black = None
        self.img_black_temp = None
        self.img_preserve = None
    
    def remove_exp_preserve(self):
        self.img = None
        self.img_temp = None
        self.img_black = None
        self.img_black_temp = None

def plt_img(img):
    plt.cla()
    plt.close('all')
    figure = plt.figure()
    canvas = FigureCanvas(figure)
    if len(img.shape) <= 2:
        plt.imshow(img, plt.cm.magma)
    else:
        plt.imshow(img)
    plt.axis('off')
    plt.title(img.shape)
    canvas.draw()
    return canvas

def add_to_graph(canvas,graph):
    x = graph.width()-20
    y = graph.height()-20
    graphicscene = QtWidgets.QGraphicsScene()
    graphicscene.addWidget(canvas).resize(x, y)
    graph.setScene(graphicscene)
    graph.repaint()

def set_range(img,low,top):
    img = img.copy()
    img[img<low] = low
    img[img>top] = top
    return img

def plt_hist(img):
    plt.cla()
    plt.close('all')
    figure = plt.figure()
    canvas = FigureCanvas(figure)
    plt.hist(img.reshape(-1,), bins=100)
    canvas.draw()
    return canvas

def read_h5ad(fname):
    adata = ad.read_h5ad(fname)
    img_data = adata.X.sum(1).reshape((5000, -1)).astype('int')
    img_data = cv2.normalize(img_data, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    return img_data

def aligning_img(our_gene,our_fluo):
    final_img = our_fluo.img_preserve
    final_img = cv2.normalize(final_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    if our_gene.flip_V:
        final_img = cv2.flip(final_img,0)
    if our_gene.flip_H:
        final_img = cv2.flip(final_img,1)
    if our_fluo.flip_V:
        final_img = cv2.flip(final_img, 0)
    if our_gene.flip_H:
        final_img = cv2.flip(final_img, 1)

    scale = (1.0/our_gene.scale_all)*our_fluo.scale_all

    print(our_fluo.height, our_gene.height,"SCALE", scale)
    final_img = cv2.resize(final_img, (0, 0), fx=scale, fy=scale)
    final_img = cv2.normalize(final_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    angle = our_fluo.angle - our_gene.angle
    final_img = rotate_with_rect(final_img, -angle)
    final_img = cv2.normalize(final_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    x = our_fluo.move_x - our_gene.move_x
    y = our_fluo.move_y - our_gene.move_y
    x = 1.0/our_gene.scale_all*x
    y = 1.0/our_gene.scale_all*y

    img = find_edge(our_gene.img_preserve,final_img,x,y)
    img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    return img

def aligning_img_ori(our_gene,our_fluo):
    final_img = our_fluo.img_preserve
    print("*"*100)
    print(final_img.shape)
    print("&"*100)
    final_img = cv2.normalize(final_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    if our_gene.flip_V:
        final_img = cv2.flip(final_img, 0)
    if our_gene.flip_H:
        final_img = cv2.flip(final_img, 1)
    if our_fluo.flip_V:
        final_img = cv2.flip(final_img, 0)
    if our_gene.flip_H:
        final_img = cv2.flip(final_img, 1)

    scale = (1.0 / our_gene.scale_all) * our_fluo.scale_all
    soft_scale = our_fluo.scale_all
    gene_scale = (1.0 / soft_scale)*our_gene.scale_all

    gene_img = np.zeros((our_gene.img_preserve.shape[0],our_gene.img_preserve.shape[1],3))
    gene_img = cv2.resize(gene_img, (0, 0), fx=gene_scale, fy=gene_scale)
    #gene_img = cv2.normalize(gene_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    angle = our_fluo.angle - our_gene.angle
    final_img = rotate_with_rect(final_img, -angle)

    final_img = cv2.normalize(final_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    x = our_fluo.move_x - our_gene.move_x
    y = our_fluo.move_y - our_gene.move_y
    x = 1.0 / soft_scale * x
    y = 1.0 / soft_scale * y

    img = find_edge(gene_img, final_img.copy(), x, y)
    this_way = img
    print("*"*100)
    print(x,y)
    print(gene_img.shape,final_img.shape)
    print("&"*100)
    #img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    return img,scale

def find_edge(img,img1,x,y):
    img = np.zeros((img.shape[0],img.shape[1],3),dtype=np.uint8)
    top, buttom, left,right = cross_area(img,img1,x,y)
    top1, buttom1, left1, right1 = cross_area(img1,img,-x,-y)

    a = buttom - top
    b = buttom1 - top1
    c = right - left
    d = right1 - left1
    f = a-b
    buttom1 = buttom1+f
    g = c-d
    right1 = right1 + g
    img[top:buttom,left:right] = img1[top1:buttom1,left1:right1]
    print(top,buttom,top1,buttom1)
    print(left,right,left1,right1)
    return img

def cross_area(img,img1,x,y):
    h = img.shape[0]
    w = img.shape[1]
    h1 = img1.shape[0]
    w1 = img1.shape[1]
    a = int(h / 2 - y - h1 / 2)
    b = int(h / 2 - y + h1 / 2)
    c = int(w / 2 + x + w1 / 2)
    d = int(w / 2 + x - w1 / 2)
    buttom = min(b, h)
    top = max(a, 0)
    right = min(c, w)
    left = max(d,0)
    return top,buttom,left,right

def set_binary(img,a=0,b=0,c=0):
    img = img.copy()
    x = np.array(img[:, :, 0])
    y = np.array(img[:, :, 1])
    z = np.array(img[:, :, 2])
    mask_0 = (x <= a) & (y <= b) & (z <= c)
    img[mask_0] = 255
    return img

def rotate_with_rect(img, degree):
    rows, cols = img.shape[0],img.shape[1]
    center = (cols / 2, rows / 2)
    M = cv2.getRotationMatrix2D(center, degree, 1)
    top_right = np.array((cols - 1, 0)) - np.array(center)
    bottom_right = np.array((cols - 1, rows - 1)) - np.array(center)
    top_right_after_rot = M[0:2, 0:2].dot(top_right)
    bottom_right_after_rot = M[0:2, 0:2].dot(bottom_right)
    new_width = max(int(abs(bottom_right_after_rot[0] * 2) + 0.5), int(abs(top_right_after_rot[0] * 2) + 0.5))
    new_height = max(int(abs(top_right_after_rot[1] * 2) + 0.5), int(abs(bottom_right_after_rot[1] * 2) + 0.5))
    offset_x = (new_width - cols) / 2
    offset_y = (new_height - rows) / 2
    M[0, 2] += offset_x
    M[1, 2] += offset_y
    dst = cv2.warpAffine(img, M, (new_width, new_height))
    #dst = set_binary(dst)
    return dst



















